// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

using System.Collections.Generic;
using System.IO;
using Microsoft.ML;
using Microsoft.ML.Runtime;
using Microsoft.ML.Runtime.Command;
using Microsoft.ML.Runtime.CommandLine;
using Microsoft.ML.Runtime.Data;
using Microsoft.ML.Runtime.Internal.Utilities;
using Microsoft.ML.Runtime.MLTesting.Inference;
using Microsoft.ML.Runtime.PipelineInference;
using Microsoft.ML.Runtime.Sweeper;

[assembly: LoadableClass(typeof(GenerateSweepCandidatesCommand), typeof(GenerateSweepCandidatesCommand.Arguments), typeof(SignatureCommand),
    "Generate Experiment Candidates", "GenerateSweepCandidates", DocName = "command/GenerateSweepCandidates.md")]

namespace Microsoft.ML.Runtime.MLTesting.Inference
{
    /// <summary>
    /// This is a command that takes as an input:
    /// 1- the schema of a dataset, in the format produced by InferSchema
    /// 2- the path to the datafile 
    /// and generates experiment candidates by combining all the transform recipes suggested for the dataset, with all the learners available for the task. 
    /// </summary>
    public sealed class GenerateSweepCandidatesCommand : ICommand
    {
        public sealed class Arguments
        {
            [Argument(ArgumentType.Required, HelpText = "Text file with data to analyze.", ShortName = "data")]
            public string DataFile;

            [Argument(ArgumentType.Required, HelpText = "Path to the folder where to store the rsp generated.", ShortName = "out")]
            public string RspOutFolder;

            [Argument(ArgumentType.AtMostOnce, HelpText = "Optional Path to the folder where the experiments will output their evaluation results. It not set it defaults to 'rspOutputFolder'. Keeping the results in the same folder makes it easy to visualize them.", ShortName = "dout")]
            public string OutputDataFolder;

            [Argument(ArgumentType.AtMostOnce, HelpText = "Optional Path to the schema definition file generated by the InferSchema command", ShortName = "schema")]
            public string SchemaDefinitionFile;

            [Argument(ArgumentType.AtMostOnce, HelpText = "Running mode: Train, TrainTest or CV. Defaults to CV if not provided.")]
            public string Mode;

            [Argument(ArgumentType.AtMostOnce, HelpText = "Text file with data for testing, if you are using the TrainTest mode.", ShortName = "test")]
            public string TestFile;

            [Argument(ArgumentType.AtMostOnce, HelpText = "Sweeper used. Options are: UniformRandom, ldrand, KDO, RandomGrid, SMAC, NM. If not provided KDO will be used.")]
            public string Sweeper;

            [Argument(ArgumentType.AtMostOnce, HelpText = "If this option is provided, the result RSP are indented and written in separate files for easier visual inspection. Otherwise all the generated RSPs are written to a single file, one RSP per line.")]
            public bool Indent;
        }

        private readonly IHost _host;
        private readonly string _dataFile;
        private readonly string _testFile;
        private readonly string _rspsOutFolder;
        private readonly string _schemaDefinitionFile;
        private readonly bool _indented;
        private readonly string _mode;
        private readonly string _sweeper;
        private readonly string _outputDataFolder;

        public GenerateSweepCandidatesCommand(IHostEnvironment env, Arguments args)
        {
            Contracts.CheckValue(env, nameof(env));
            _host = env.Register("GenerateCandidates");
            _host.CheckValue(args, nameof(args));

            var files = new MultiFileSource(args.DataFile);
            _host.CheckUserArg(files.Count > 0, nameof(args.DataFile), "dataFile is required");
            _dataFile = args.DataFile;

            _rspsOutFolder = Utils.CreateFolderIfNotExists(args.RspOutFolder);
            _host.CheckUserArg(_rspsOutFolder != null, nameof(args.RspOutFolder), "Provide a value rspOutFolder (or 'out', the short name).");

            if (!string.IsNullOrWhiteSpace(args.SchemaDefinitionFile))
            {
                Utils.CheckOptionalUserDirectory(args.SchemaDefinitionFile, nameof(args.SchemaDefinitionFile));
                _schemaDefinitionFile = args.SchemaDefinitionFile;
            }

            if (!string.IsNullOrWhiteSpace(args.Sweeper))
            {
                var info = ComponentCatalog.GetLoadableClassInfo<SignatureSweeper>(args.Sweeper);
                _host.CheckUserArg(info?.SignatureTypes[0] == typeof(SignatureSweeper), nameof(args.Sweeper),
                    "Please specify a valid sweeper.");
                _sweeper = args.Sweeper;
            }
            else
                _sweeper = "kdo";

            if (!string.IsNullOrWhiteSpace(args.Mode))
            {
                var info = ComponentCatalog.GetLoadableClassInfo<SignatureCommand>(args.Mode);
                _host.CheckUserArg(info?.Type == typeof(TrainCommand) ||
                                   info?.Type == typeof(TrainTestCommand) ||
                                   info?.Type == typeof(CrossValidationCommand), nameof(args.Mode), "Invalid mode.");
                _mode = args.Mode;
            }
            else
                _mode = CrossValidationCommand.LoadName;

            _indented = args.Indent;

            if (!string.IsNullOrWhiteSpace(args.TestFile))
            {
                files = new MultiFileSource(args.TestFile);
                _host.CheckUserArg(files.Count > 0, nameof(args.TestFile), "testFile needs to be a valid file, if provided.");
                _testFile = args.TestFile;
            }
            else
            {
                _host.CheckUserArg(_mode != TrainTestCommand.LoadName, nameof(args.TestFile), "testFile needs to be a valid file, for mode = TrainTest.");
            }

            _outputDataFolder = Utils.CreateFolderIfNotExists(args.OutputDataFolder);
            if (_outputDataFolder == null)
                _outputDataFolder = _rspsOutFolder;
        }

        public void Run()
        {
            using (var ch = _host.Start("Running"))
            {
                RunCore(ch);
                ch.Done();
            }
        }

        private void RunCore(IChannel ch)
        {
            _host.AssertValue(ch);

            var candidates = ExperimentsGenerator.GenerateCandidates(_host, _dataFile, _schemaDefinitionFile);

            // If the user wants the results indented, print every rsp created in its own file,  indented for easy inspection.
            if (_indented)
                Print(candidates, ch);

            // Otheriwse, print all RSPs in one single file, one rsp per line.
            else
            {
                using (var sw = new StreamWriter(Path.Combine($"{_rspsOutFolder}", "Experiments.rsp")))
                    Print(candidates, ch, sw);
            }
        }

        private void Print(List<ExperimentsGenerator.Sweep> candidates, IChannel ch, StreamWriter sw = null)
        {
            for (int i = 0; i < candidates.Count; i++)
            {
                var experiment = candidates[i];

                string rsp = null;
                // If there is a sweeper section, build a sweep command, otherwise just format the pattern to rsp.
                if (experiment.TrainerSweeper.Parameters.Count > 0)
                    rsp = $" {SweepCommand.LoadName} evaluator=Local{{ pattern= {{{experiment.Pattern.ToStringRep(_mode, _dataFile, _testFile)}}} outfolder={_outputDataFolder} }} {experiment.TrainerSweeper.ToStringRep(_sweeper)}";
                else
                    rsp = experiment.Pattern.ToStringRep(_mode, _dataFile, _testFile, Path.Combine(_outputDataFolder, $"{i}_{experiment.Pattern.Learner.LoadableClassInfo.LoadNames[0]}.out"));

                ch.Info(rsp);

                if (sw != null)
                    sw.WriteLine(rsp);
                else
                {
                    // Print everything in its file indented.
                    using (var swr = new StreamWriter(Path.Combine(_rspsOutFolder, $"Experiment_{i:0000}.rsp")))
                        swr.WriteLine(CmdIndenter.GetIndentedCommandLine(rsp));
                }
            }
        }
    }
}
